# -*- coding:utf-8 -*-
import carla
import math
import numpy as np
import matplotlib.pyplot as plt
from my_controller.controller import Vehicle_control
from planner import planning_utils
from my_utils import generate_dlc_path

class Vehicle:
    def __init__(self, world, spawn_point_index):
        self.world = world
        self.spawn_point = world.get_map().get_spawn_points()[spawn_point_index]
        self.model = self.spawn_vehicle()

    def spawn_vehicle(self):
        blueprint = self.world.get_blueprint_library().find('vehicle.tesla.model3')
        blueprint.set_attribute('color', '255,88,0')
        vehicle = self.world.spawn_actor(blueprint, self.spawn_point)
        physics_control = carla.VehiclePhysicsControl()
        physics_control.mass = 1412
        vehicle.apply_physics_control(physics_control)
        return vehicle

    def get_vehicle_info(self):
        physics_control = self.model.get_physics_control()
        wheel1 = physics_control.wheels[1].position
        wheel3 = physics_control.wheels[3].position
        dis13 = math.sqrt((wheel1.x - wheel3.x) ** 2 + (wheel1.y - wheel3.y) ** 2)
        wheelbase = dis13 / 100
        lf = wheelbase / 2
        lr = wheelbase / 2
        mass = physics_control.mass
        return (lr, lf, -110000, -110000, mass, 1537)

    def apply_control(self, control):
        self.model.apply_control(control)

    def get_transform(self):
        return self.model.get_transform()

    def destroy(self):
        self.model.destroy()

class Controller:
    def __init__(self, vehicle, path, controller_type="MPC_controller"):
        self.vehicle = vehicle
        self.path = path
        self.controller = Vehicle_control(vehicle.model, vehicle.get_vehicle_info(), path, controller_type)

    def run_step(self, target_speed):
        control = self.controller.run_step(target_speed)
        return control

class World:
    def __init__(self, client_address, world_name):
        self.client = carla.Client(*client_address)
        self.client.set_timeout(10)
        self.world = self.client.load_world(world_name)

    def tick(self):
        self.world.tick()

    def get_spectator(self):
        return self.world.get_spectator()

    def set_spectator(self, transform):
        self.get_spectator().set_transform(transform)

    def get_settings(self):
        return self.world.get_settings()

    def set_settings(self, settings):
        self.world.apply_settings(settings)

class Simulation:
    def __init__(self, world, vehicle, controller, draw):
        self.world = world
        self.vehicle = vehicle
        self.controller = controller
        self.draw = draw

    def run(self):
        try:
            max_speed = 50
            while True:
                self.world.tick()
                spectator_transform = carla.Transform(self.vehicle.get_transform().location + carla.Location(z=50), carla.Rotation(pitch=-90))
                self.world.set_spectator(spectator_transform)
                control = self.controller.run_step(max_speed)
                self.vehicle.apply_control(control)
                self.draw.update(max_speed)
                if self.draw.get_distance(self.controller.path) < 2:
                    self.vehicle.apply_control(carla.VehicleControl(steer=0, throttle=0, brake=1))
                    break
        finally:
            self.vehicle.destroy()
            settings = self.world.get_settings()
            settings.synchronous_mode = False
            settings.fixed_delta_seconds = None
            self.world.set_settings(settings)

class Draw:
    def __init__(self, vehicle):
        self.vehicle = vehicle
        self.data = {
            'direction': [],
            'speed': [],
            'target_speed': [],
            'lat_acc': [],
            'car_position': [],
            'pathway': []
        }

    def update(self, max_speed):
        self.data['lat_acc'].append(self.get_lat_acc())
        self.data['direction'].append(self.get_direction())
        self.data['speed'].append(self.get_velocity())
        self.data['target_speed'].append(max_speed)
        location = self.vehicle.get_transform().location
        self.data['car_position'].append([location.x, location.y])

    def get_lat_acc(self):
        acc = self.vehicle.model.get_acceleration()
        acc_vec = np.array([acc.x, acc.y, acc.z])
        right_vec = self.vehicle.get_transform().get_right_vector()
        right_vec = np.array([right_vec.x, right_vec.y, right_vec.z])
        lat_acc = np.dot(acc_vec, right_vec)
        return lat_acc

    def get_velocity(self):
        velocity = self.vehicle.model.get_velocity()
        return 3.6 * math.sqrt(velocity.x ** 2 + velocity.y ** 2 + velocity.z ** 2)

    def get_direction(self):
        return self.vehicle.get_transform().rotation.yaw * (math.pi / 180)

    def get_distance(self, path):
        location = self.vehicle.get_transform().location
        return location.distance(path[-1][0].transform.location)

    def show(self):
        plt.subplot(2, 2, 1)
        plt.plot(self.data['direction'])
        plt.ylim(bottom=-5, top=5)
        plt.title('delta')

        plt.subplot(2, 2, 2)
        plt.plot(self.data['speed'])
        plt.plot(self.data['target_speed'], color="r")
        plt.ylim(bottom=0, top=max(self.data['target_speed']) + 10)
        plt.title('speed')

        plt.subplot(2, 2, 3)
        x1 = [point[0] for point in self.data['car_position']]
        y1 = [point[1] for point in self.data['car_position']]
        plt.plot(x1, y1, label='real', marker='o', alpha=0.5, markersize=1, linestyle='-', color='blue')
        x2 = [point.x for point in self.data['pathway']]
        y2 = [point.y for point in self.data['pathway']]
        plt.plot(x2, y2, label='refer', marker='s', alpha=0.5, markersize=1, linestyle='-', color='orange')
        plt.legend()
        plt.title("path")

        plt.subplot(2, 2, 4)
        plt.plot(self.data['lat_acc'])
        plt.title('l_acc')
        plt.show()

def main():
    client_address = ("localhost", 2000)
    world_name = 'town05'
    spawn_point_index = 7

    world = World(client_address, world_name)
    vehicle = Vehicle(world.world, spawn_point_index)
    path = generate_dlc_path.get_waypoints_2(vehicle.spawn_point)
    controller = Controller(vehicle, path)
    draw = Draw(vehicle)
    draw.pathway = [obj[0].transform.location for obj in path]
    simulation = Simulation(world, vehicle, controller, draw)
    simulation.run()
    draw.show()

if __name__ == "__main__":
    main()
